"""
容器管理
"""
from .base import docker_to_dict
from typing import Literal, Any, Dict, List, Optional
from pydantic import Field
from mcp.server import FastMCP
from .docker_settings import DockerSettings
from docker.models.containers import Container
import docker


def get_container(
        _docker: docker.DockerClient,
        name: str = Field(default="python-sandbox", description="容器名称."),
        container_id: Optional[str] = Field(default=None, description="容器id."),
) -> Optional[Container]:
    """根据容器名称/容器id找到容器"""
    if container_id is not None:
        container = _docker.containers.get(container_id)
        return container
    else:
        containers = _docker.containers.list(all=True, filters={
            'name': name
        })
        for container in containers:
            if container.name == name:
                return container
        return None


def add_container_tools(mcp: FastMCP, _docker: docker.DockerClient):
    """
    仅添加跟python沙箱相关的接口
    """
    @mcp.tool()
    async def list_python_containers(
            status: Literal["restarting", "running", "paused", "exited"] = Field(default="running", description="列出指定运行状态的python容器."),
            name: str = Field(default="python-sandbox", description="容器名称.")
    ) -> List[Dict]:
        """
            列出所有的python沙箱。
            Returns:
                返回python沙箱的容器列表。
            Raises:
                如果报错，将抛出`docker.errors.APIError`类型的异常.
        """
        containers = _docker.containers.list(all=True, filters={
            'status': status,
            'name': name
        })
        result = [docker_to_dict(c) for c in containers]
        return result

    @mcp.tool()
    async def get_python_container(
            name: str = Field(default="python-sandbox", description="容器名称."),
            container_id: Optional[str] = Field(default=None, description="容器id."),
    ) -> Optional[Dict[str, Any]]:
        """
        获取指定名称或id的python沙箱
        Returns:
            返回python沙箱的详细信息. 如果未找到则返回None
        """
        container = get_container(_docker, name, container_id)
        if container is not None:
            container = docker_to_dict(container)
        return container

    @mcp.tool()
    async def create_python_container(
            name: str = Field(default="python-sandbox", description="容器名称."),
            python_version: str = Field(default="python:3.12", description="python版本号."),
            requirements_file: Optional[str] = Field(default=None, description="python环境依赖包文件."),
            files_dir: str = Field(default="/files", description="文件存放目录."),
            port: int = Field(default=8194, description="容器的暴露端口.")
    ):
        """创建一个python沙箱。
            Returns:
                如果已存在name的沙箱，则直接返回该沙箱的信息，否则创建新的python沙箱并运行，然后返回沙箱的信息。
            Raises:
                如果报错，将抛出`docker.errors.ImageNotFound`类型的异常或 `docker.errors.APIError`的异常。
        """

        # 检查是否存在该名称的docker, 如果存在，则直接返回。
        container = get_container(_docker, name=name, container_id=None)
        if container is not None:
            return docker_to_dict(container)

        # 配置 Dockerfile
        settings = DockerSettings()
        settings.set_python_version(python_version)
        if requirements_file is not None:
            settings.set_requirements(requirements_file)
        settings.set_file_dir(files_dir)
        settings.set_port(port)
        settings.save()

        # 从当前sandboxes的python目录文件创建一个镜像
        args = {
            'path': settings.sandbox_dir.as_posix(),
            'tag': name
        }
        image, _ = _docker.images.build(**args)
        image = docker_to_dict(image)

        # 从镜像创建并运行容器
        args = {
            'image': image['id'],
            'name': f'{name}',
            'network': 'bridge',
            'ports': {'8194': port}
        }
        container = _docker.containers.create(**args)
        container.start()
        return docker_to_dict(container)

    @mcp.tool()
    async def run_python_container(
            name: str = Field(default="python-sandbox", description="容器名称."),
            container_id: Optional[str] = Field(default=None, description="容器id."),
    ) -> Optional[Dict[str, Any]]:
        """运行容器
            Returns:
                返回当前运行的python沙箱的信息。
            Raises:
                如果报错，将抛出`docker.errors.ImageNotFound`类型的异常或 `docker.errors.APIError`的异常。
        """
        container = get_container(_docker, name, container_id)
        if container is None:
            return None
        if container.status != 'running':
            container.start()
        container = get_container(_docker, name, container_id)
        return docker_to_dict(container)

    @mcp.tool()
    async def stop_python_container(
            name: str = Field(default="python-sandbox", description="容器名称."),
            container_id: Optional[str] = Field(default=None, description="容器id."),
    ) -> Optional[Dict[str, Any]]:
        """停止运行中的python沙箱"""
        container = get_container(_docker, name, container_id)
        if container is None:
            return None
        if container.status == 'running':
            container.stop()
        container = get_container(_docker, name, container_id)
        return docker_to_dict(container)

    @mcp.tool()
    async def fetch_python_container_logs(
            name: str = Field(default="python-sandbox", description="容器名称."),
            container_id: Optional[str] = Field(default=None, description="容器id."),
            tail: int | Literal["all"] = Field(100, description="日志文件末尾的行数.")
    ) -> Optional[List[str]]:
        """获取python沙箱的运行日志"""
        container = get_container(_docker, name, container_id)
        if container is None:
            return None

        logs = container.logs(tail=tail).decode("utf-8")
        return logs.split("\n")

    @mcp.tool()
    async def remove_python_container(
            name: str = Field(default="python-sandbox", description="容器名称."),
            container_id: Optional[str] = Field(default=None, description="容器id."),
            force: bool = Field(False, description='是否使用SIGKILL强制移除运行中的容器.')
    ) -> Optional[Dict[str, Any]]:
        """删除python沙箱. """
        container = get_container(_docker, name, container_id)
        if container is None:
            return None
        if container.status == "running":
            container.stop()
        container.remove(force=force)
        return docker_to_dict(container, {"status": "removed"})
